import cv2
import os

# 用户可配置参数
video_path = "/home/shuai/Videos/循迹避障数据集.mp4"  # 替换为您的MP4文件路径
output_dir = "/home/shuai/Videos/output_frames"   # 替换为您想保存帧的目录
frames_per_second = 3         # 每秒抽取的帧数 (x值)

# 创建输出目录
os.makedirs(output_dir, exist_ok=True)

# 打开视频文件
cap = cv2.VideoCapture(video_path)
if not cap.isOpened():
    print("Error: 无法打开视频文件")
    exit()

# 获取视频信息
fps = cap.get(cv2.CAP_PROP_FPS)
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
duration = total_frames / fps
frame_interval = int(round(fps / frames_per_second))

print(f"视频信息: {fps:.2f} FPS, 总帧数: {total_frames}, 时长: {duration:.2f}秒")
print(f"将以每秒 {frames_per_second} 帧的速度抽取，每 {frame_interval} 帧抽取一次")

# 开始抽帧
count = 0
frame_number = 0
saved_count = 0

while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break
    
    # 按间隔抽取帧
    if frame_number % frame_interval == 0:
        # 格式化文件名 (frame_0000.jpg, frame_0001.jpg, ...)
        filename = f"frame_{saved_count:04d}.jpg"
        output_path = os.path.join(output_dir, filename)
        
        # 保存帧为JPEG图像
        cv2.imwrite(output_path, frame)
        saved_count += 1
    
    frame_number += 1
    count += 1

    # 显示进度
    if count % 100 == 0:
        print(f"处理进度: {count}/{total_frames} 帧", end="\r")

# 释放资源
cap.release()
print("\n完成!")
print(f"共保存了 {saved_count} 帧图像到目录: {output_dir}")